import numpy as np
import cv2
import open3d as o3d
import sys
import yaml
from pathlib import Path

from intrinsics import intrinsics_sim, intrinsics_real, width_sim, height_sim
from cam_yaml_new import create_transformation_matrix_from_plf



def project_points_to_image(points, colors, camera_matrix, intrinsics):
    fx, fy, cx, cy = intrinsics
    points_cam = camera_matrix @ np.column_stack([points, np.ones(len(points))]).T
    points_cam = points_cam[:3, :].T
    coord_convert_inv = np.array([[1, 0, 0], [0, -1, 0], [0, 0, -1]])
    points_cam = points_cam @ coord_convert_inv.T
    valid_mask = points_cam[:, 2] > 0
    points_cam = points_cam[valid_mask]
    colors = colors[valid_mask]
    u = (points_cam[:, 0] * fx / points_cam[:, 2] + cx).astype(int)
    v = (points_cam[:, 1] * fy / points_cam[:, 2] + cy).astype(int)
    valid_pixel_mask = (u >= 0) & (u < width_sim) & (v >= 0) & (v < height_sim)
    u = u[valid_pixel_mask]
    v = v[valid_pixel_mask]
    colors = colors[valid_pixel_mask]
    return u, v, colors

def load_camera_pose_from_yaml(yaml_path, camera_name):
    with open(yaml_path, 'r', encoding='utf-8') as f:
        yaml_data = yaml.safe_load(f)
    camera_list = yaml_data['static_camera_list']
    for cam in camera_list:
        if cam['name'] == camera_name:
            return {
                'position': np.array(cam['position']),
                'forward': np.array(cam['forward']),
                'left': np.array(cam['left'])
            }
    raise ValueError(f"Camera '{camera_name}' not found in YAML file")

def load_pcd_as_numpy(pcd_path):
    pcd = o3d.io.read_point_cloud(str(pcd_path))
    points = np.asarray(pcd.points)
    colors = np.asarray(pcd.colors)
    if len(colors) == 0:
        colors = np.ones((len(points), 3)) * 0.5
    return np.column_stack([points, colors])

def render_pcd_image(pcd_data, camera_position, camera_forward, camera_left, intrinsics):
    points = pcd_data[:, :3]
    colors = pcd_data[:, 3:6]
    if colors.max() > 1.0:
        colors = colors / 255.0
    camera_matrix = np.linalg.inv(create_transformation_matrix_from_plf(camera_position, camera_left, camera_forward))
    u, v, colors = project_points_to_image(points, colors, camera_matrix, intrinsics)
    image = np.zeros((height_sim, width_sim, 3), dtype=np.uint8)
    if len(u) > 0:
        image[v, u] = (colors * 255).astype(np.uint8)
    return image

def visualize_camera_view(yaml_path, pcd_path, camera_name, output_path=None):
    camera_pose = load_camera_pose_from_yaml(yaml_path,camera_name)
    pcd_data = load_pcd_as_numpy(pcd_path)
    image = render_pcd_image(pcd_data, camera_pose['position'], camera_pose['forward'], camera_pose['left'], intrinsics_sim)
    print(f"image shape: {image.shape}")
    if output_path:
        cv2.imwrite(str(output_path), image)
        print(f"Camera view saved to {output_path}")
    return image
